#!/usr/bin/python
#Author: Jimmy Saw
#Usage: python get_lineage.py "Genus species"
#It is supposed to give you the parent phylum of the query term

# NEED TO FIX!!!!

import sys
import re
from Bio import Entrez

Entrez.email = "jimmysaw@gmail.com"

f = open(sys.argv[1], "rU")

lines = f.readlines()

num = len(lines)

#rx = re.compile('(\w+ \w+).*') #only keep genus and species names as search query
rx = re.compile('(\w+).*') #only keep genus name as search query

m1 = re.compile('(.*\[.*\].*)\[(.*)\]')
m2 = re.compile('(.*)\[(.*)\]')

n = 0

while n < num:
    curr_line = lines[n].split('\t')
    curr_prodigal_id = curr_line[0]
    curr_locus_tag = curr_line[1]
    curr_feat_type = curr_line[2]
    curr_start = int(curr_line[3])
    curr_stop = int(curr_line[4])
    curr_frame = curr_line[5]
    curr_blast_status = curr_line[6]
    curr_blast_evalue = curr_line[7]
    curr_tophit_id = curr_line[8]
    curr_tophit_sim = curr_line[9]
    curr_tophit_acc = curr_line[10]
    curr_tophit_gi = curr_line[11]
    curr_hit_def = curr_line[17]
    curr_hit_org = curr_line[18]
    curr_cluster_def = curr_line[31]
    curr_cluster_org = curr_line[32]

    if curr_blast_status == "Yes":
        blast_file = curr_prodigal_id + ".blastp.tbl"
        bf = open(blast_file, "rU")
        bfl = bf.readlines()
        ln = len(bfl)
        #Assignment of weights to decide taxonomic affiliation
        if ln == 1:
            searchterm = curr_hit_org
        elif 1 < ln < 10:
            c1 = bfl[0].split('\t')
            s1 = float(c1[0]) #best hit score
            i = 0
            weight_sum = 0.0
            weight_list = []
            org_list = []
            #mydict = {}
            while i < ln:
                cl = bfl[i].split('\t')
                score = float(cl[0])
                weight = (score/s1)**32
                org = cl[13].rstrip()
                if m1.match(org):
                    pattern = m1.match(org)
                elif m2.match(org):
                    pattern = m2.match(org)
                else:
                    organism = ""
                if organism != "":
                    organism = pattern.group(2)
                else:
                    organism = ""
                #mydict[organism] = weight
                weight_sum += weight
                weight_list.append(weight)
                org_list.append(organism)
                i += 1
            for i, j in enumerate(weight_list):
                threshold = 0.75 * weight_sum
                if j >= threshold:
                    searchterm = org_list[i]
                else:
                    searchterm = ""
        else:
            c1 = bfl[0].split('\t')
            s1 = float(c1[0]) #best hit score
            i = 0
            weight_sum = 0.0
            weight_list = []
            org_list = []
            while i < 10:
                cl = bfl[i].split('\t')
                score = float(cl[0])
                weight = (score/s1)**32
                org = cl[13].rstrip()
                if m1.match(org):
                    pattern = m1.match(org)
                elif m2.match(org):
                    pattern = m2.match(org)
                else:
                    organism = ""
                if organism != "":
                    organism = pattern.group(2)
                else:
                    organism = ""
                weight_sum += weight
                weight_list.append(weight)
                org_list.append(organism)
                i += 1
            for i, j in enumerate(weight_list):
                threshold = 0.75 * weight_sum
                if j >= threshold:
                    searchterm = org_list[i]
                else:
                    searchterm = ""

        #searchterm = curr_hit_org
        if searchterm != "":
            pattern = rx.match(searchterm)
            st = pattern.group(1)
            handle = Entrez.esearch(db="Taxonomy", term=st)
            record = Entrez.read(handle)
            taxid = record["IdList"]
            #print curr_prodigal_id, "\t", taxid[0]
            if len(taxid) != 0:
                h = Entrez.efetch(db="Taxonomy", id=taxid, retmode="xml")
                rec = Entrez.parse(h)
                for r in rec:
                    lineage = r.get("Lineage")
                    l = lineage.split(';')
                    print curr_prodigal_id, "\t", l[2].lstrip()

    n += 1

